#!/usr/bin/env python3
# -*- coding: utf-8 -*-

''' 
    date: 2021-03-18
    author: baole 

    从东财爬取总股本和流通股本
    原网址：http://f10.eastmoney.com/f10_v2/CapitalStockStructure.aspx?code=sz300171
    里面的单位是万股。

    遍历所有的股票代码，获取股本结构中的股本总数和已上市流通A股。
    保存到数据库中 序号，股票代码，总股本，流通股本（A股）
'''

import os
import sys
import requests
import json
import pymysql
import time
import random


from quantway.lib.utils import read_jsonparam, get_module_name, get_codelist, split_list_by_num





def get_guben(confpath, li_stock_code):
    '''
        从东财爬取总股本和流通股本
    '''

    # 获取参数
    param = read_jsonparam(confpath)
    db_param = param.database

    try:
        db = pymysql.connect(**db_param)
    except Exception as e:
        print("[%s]连接数据库失败,程序终止" % (get_module_name()), e)
        return
    else:
        print('[%s]:数据库连接成功' % get_module_name())


    cursor = db.cursor()
    li_fail = list()            # 失败的股票列表
    num_total = len(li_stock_code)      # 本进程处理的股票总数
    num_fail = 0                        # 失败数量

    for i, stock_code in enumerate(li_stock_code):
        # 随机休息1~20s， 否则容易被东财屏蔽ip一段时间
        # time.sleep(random.randint(1,20))
        
        # 抓取数据
        r = query_guben_from_eastmoney(stock_code.upper())

        
        if r is None:
            num_fail += 1 
            li_fail.append(stock_code)
            print("[%d/%d] 获取%s的数据失败" % (i+i, num_fail, stock_code))
            continue
        stock_share = json.loads(r.text)
        if stock_share == None or len(stock_share) == 0:
            num_fail += 1 
            li_fail.append(stock_code)
            print("[%d/%d] 获取%s的数据失败" % (i+i, num_fail, stock_code))
            continue 
        if 'status' in stock_share and stock_share['status'] == -1:
            num_fail += 1 
            li_fail.append(stock_code)
            print("[%d/%d] 获取%s的数据失败 %s" % (i+i, num_fail, stock_code, stock_share['message']))
            continue

        try:
            # 总股本
            zgb = stock_share['CapitalStockStructureDetail']["zgb"]
            zgb = float(zgb.replace(",",""))
            # 以上市流通A股
            yssltag = stock_share['CapitalStockStructureDetail']["yssltag"]
            yssltag = float(yssltag.replace(",",""))
        except Exception as e:
            print("%s convertion error" % stock_code, e.args)
            num_fail += 1  
            li_fail.append(stock_code)
            print("[%d/%d] 获取%s的数据失败" % (i+i, num_fail, stock_code))
            continue
        
        try:
            db.begin()
            # 已经存在，则比较，如果有差异就更新。
            sql = "select zgb, yssltag from guben where stock_code='%s'" % stock_code
            if cursor.execute(sql) > 0:
                row = cursor.fetchone()
                if row[0] == zgb and row[1] == yssltag:
                    print("[%d/%d] %s的股本和流通A股没有变化" % (i+1, num_total, stock_code ))
                else:
                    # 更新总股本和流通A股
                    sql = "update guben set zgb=%.4f and yssltag=%.4f where stock_code='%s'" % (zgb, yssltag, stock_code)
                    cursor.execute(sql)
            else:
                sql = "insert into guben (stock_code, zgb, yssltag, lt5cgs, sjltag) values('%s', %.4f, %.4f, 0, 0)" % (stock_code, zgb, yssltag)
                cursor.execute(sql)
            db.commit()
            print("[%d/%d] 获得股票%s的股本 " % (i+1, num_total, stock_code ))
        except Exception as e:
            db.rollback()
            print("[%d/%d]%s保存数据失败" % stock_code, e)
            num_total += 1 
            li_fail.append(stock_code)


    cursor.close()
    db.close()

    print("总共:%d股票，失败%d股票" % (num_total, num_fail))
    if num_fail > 0:
        print("失败的股票：", ",".join(li_fail))



def query_guben_from_eastmoney(stock_code):
    # 对于失败的request,我们会尝试5次。
    attempt = 0
    while attempt<5:
        try:
            r = requests.get('http://f10.eastmoney.com/CapitalStockStructure/CapitalStockStructureAjax?code=%s' % stock_code)
            return r                
        except Exception as e:
            print(e)
            time.sleep(60)
        attempt+=1
        print("获取%s的股本超时，重试%d " % (stock_code, attempt) )


    return None


        


    


if __name__ == '__main__':

    li_stock_code = get_codelist("all_stockcode")    

    if len(sys.argv) > 2:
        total = int(sys.argv[1])
        sub = int(sys.argv[2])
        split = split_list_by_num(li_stock_code, total)
        get_guben("conf/spider.json", split[sub])
    else:
        get_guben("conf/spider.json", li_stock_code)

    
